# Copyright (c) 2024-present AI-Labs

from fastapi import APIRouter, Request, HTTPException, Response, UploadFile

import torch

import os, sys
import uuid
from datetime import datetime

from ipex_llm.transformers.npu_model import FunAsrAutoModel as AutoModel
from funasr.utils.postprocess_utils import rich_transcription_postprocess

from utils.FileUtil import fetch_file
from configs import config


# 定义路由信息
router = APIRouter(
    prefix='/asr/funasr',
    tags = ['语音识别']
)


# 加载语音识别模型
model = AutoModel(
    model=config.service.funasr.model_path,
    attn_implementation="eager",
    load_in_low_bit=config.service.funasr.low_bit,
    low_cpu_mem_usage=True,
    optimize_model=True,
    save_directory=config.service.funasr.save_directory,
)


"""
语音识别具体实现逻辑
"""
def audio_asr(source):
    # 接收前端请求的语音数据并保存到本地
    localdir = f"upload/{datetime.now().strftime('%Y-%m-%d')}"
    os.makedirs(f"{config.setting.statics.path}/{localdir}", exist_ok=True)
    localfile = f"{localdir}/{uuid.uuid4()}"

    # 支持Base64编码的语音数据和语音文件的URL地址
    fetch_file(source, f"{config.setting.statics.path}/{localfile}")

    # 模型推理，进行语音数据识别
    res = model.generate(
        input=f"{config.setting.statics.path}/{localfile}",
        batch_size_s=300,
        hotword='魔搭',
    )

    # 获取语音识别结果并返回
    result = rich_transcription_postprocess(res[0]["text"].replace(" ",""))

    return result


"""
对外暴露的语音识别接口，接收Base64编码的语音数据
"""
@router.post("/audio/translations")
def audio_translations(request: dict):
    return {"text": audio_asr(request["file"])}


"""
对外暴露的语音识别接口，接收上传的语音数据二进制文件
"""
@router.post("/audio_to_text")
def audio_to_text(audio: UploadFile):
    # 接收前端请求的语音数据并保存到本地
    localdir = f"statics/upload/{datetime.now().strftime('%Y-%m-%d')}/{uuid.uuid4()}"
    os.makedirs(localdir, exist_ok=True)
    localfile = f"{localdir}/{audio.filename}"

    # 将二进制的语音文件保存到本地
    with open(localfile, "wb") as f:
        f.write(audio.file.read())

    # 模型推理，进行语音数据识别
    res = model.generate(
        input=localfile,
        batch_size_s=300,
        hotword='魔搭',
    )

    # 获取语音识别结果并返回
    result = rich_transcription_postprocess(res[0]["text"].replace(" ",""))

    return result
